import matplotlib.pyplot as plt  # 导入绘图库，用于数据可视化
import numpy as np  # 导入数组库，用于生成S形曲线
from sklearn.datasets import make_s_curve  # 导入Scikit-learn中的make_s_curve函数，生成3D S形曲线
import torch  # 导入PyTorch包，用于深度学习开发
import io
from PIL import Image  # 导入Python Imaging Library，用于图像处理
from loss import diffusion_loss_fn
from model import MLPDiffusion
from sample import p_sample_loop

# 生成一万个点，得到s curve
s_curve, _ = make_s_curve(10 ** 4, noise=0.1)  # 生成10000个3D数据点，noise参数表示给数据加入噪声
s_curve = s_curve[:, [0, 2]] / 10.0  # 只保留第一个和第三个坐标，将数据缩小到1/10，处理成2D数据
# s_curve[:,[0,2]]：选取第一个和第三个坐标
# /10.0将数据缩小到1/10，处理成2D数据

print("shape of s:", np.shape(s_curve))

data = s_curve.T  # 将数据转置
fig, ax = plt.subplots()  # 创建子图
ax.scatter(*data, color='blue', edgecolor='white')  # 绘制散点图
ax.axis('off')
dataset = torch.Tensor(s_curve).float()
# 画出 s curve
plt.show()

# 确定超参数的值
num_steps = 100  # 可由beta值估算

# 制定每一步的beta，beta按照时间从小到大变化
betas = torch.linspace(-6, 6, num_steps)  # 在-6到6之间生成100个数
betas = torch.sigmoid(betas) * (0.5e-2 - 1e-5) + 1e-5  # 对生成的100个数进行sigmoid归一化和缩放
# sigmoid归一化将数据缩放到(0, 1)之间，缩放的目的是将值域分布到较小的区间内，从而便于计算机进行计算

# 计算alpha、alpha_prod、alpha_prod_previous、alpha_bar_sqrt等变量的值
alphas = 1 - betas
# alpha连乘
alphas_prod = torch.cumprod(alphas, 0)  # 累乘函数，计算从t=0到当前时刻的所有alpha值的积
alphas_prod_p = torch.cat([torch.tensor([1]).float(), alphas_prod[:-1]], 0)  # 求导需要用到前一时刻的alphas_prod，第一项需要设为1，然后进行拼接
# alphas_prod开根号
alphas_bar_sqrt = torch.sqrt(alphas_prod)  # 将alphas_prod全部开根
# 之后公式中要用的
one_minus_alphas_bar_log = torch.log(1 - alphas_prod)
one_minus_alphas_bar_sqrt = torch.sqrt(1 - alphas_prod)

# 大小都一样，常数不需要训练
assert alphas.shape == alphas_prod.shape == alphas_prod_p.shape == \
       alphas_bar_sqrt.shape == one_minus_alphas_bar_log.shape \
       == one_minus_alphas_bar_sqrt.shape  # 断言alphas等变量的大小都相同。
print("all the same shape", betas.shape)


# 给定初始，算出任意时刻采样值——正向扩散
# 计算任意时刻的x采样值，基于x_0和重参数化
def forawrd_diffusion(x_0, t):
    """可以基于x[0]得到任意时刻t的x[t]"""
    # 生成正态分布采样
    noise = torch.randn_like(x_0)  # 生成与x_0同样shape的随机标准正态分布的采样值
    # 得到均值方差
    alphas_t = alphas_bar_sqrt[t]  # 获取alpha_bar_sqrt中第t个位置的值
    alphas_1_m_t = one_minus_alphas_bar_sqrt[t]  # 获取one_minus_alphas_bar_sqrt中第t个位置的值
    # 根据x0求xt
    return (alphas_t * x_0 + alphas_1_m_t * noise)  # 在x[0]的基础上添加噪声


# 演示加噪过程，加噪100步情况
num_shows = 20  # 设置展示20个步骤
fig, axs = plt.subplots(2, 10, figsize=(28, 3))  # 创建子图，2行10列
plt.rc('text', color='black')  # 设置图片字体的颜色为黑色

# 共有10000个点，每个点包含两个坐标
# 生成100步以内每隔5步加噪声后的图像，扩散过程散点图演示——基于x0生成条件分布采样得到xt
for i in range(num_shows):  # 遍历每个展示步骤
    j = i // 10  # 获得横轴索引值，保证不会超过两行
    k = i % 10  # 获得列索引值
    q_i = forawrd_diffusion(dataset, torch.tensor([i * num_steps // num_shows]))  # 数据集传入模型中，获得任意时刻的x采样值
    axs[j, k].scatter(q_i[:, 0], q_i[:, 1], color='red', edgecolor='white')  # 绘制散点图
    axs[j, k].set_axis_off()  # 隐藏轴线
    axs[j, k].set_title('$q(\mathbf{x}_{' + str(i * num_steps // num_shows) + '})$')
plt.show()  # 显示图片


# 逆向阶段训练、采样过程
# 开始训练模型，打印loss以及中间的重构效果
seed = 1234

print('Training model...')
batch_size = 128
# dataset放到dataloader中
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)
# 迭代周期
num_epoch = 4000
plt.rc('text', color='blue')  # 设置图形的文本渲染方式为TeX语法渲染
# 实例化模型,传入一个数
model = MLPDiffusion(num_steps)  # 输出维度是2,输入是x和step
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
# epoch遍历
for t in range(num_epoch):
    # dataloader遍历，进行批量训练和梯度下降
    for idx, batch_x in enumerate(dataloader):
        # 计算损失loss
        loss = diffusion_loss_fn(model, batch_x, alphas_bar_sqrt, one_minus_alphas_bar_sqrt, num_steps)
        optimizer.zero_grad()  # 将梯度归零，计算每个参数的梯度值
        loss.backward()  # 反向传播计算梯度值
        # 梯度clip, 将梯度约束在某一个区间之内,在优化器更新之前进行梯度截断操作,避免梯度爆炸
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.)
        optimizer.step()  # 执行一步参数更新，对损失进行反向传播更新
    # 每100步进行采样、打印生成效果
    if t % 100 == 0:
        print(loss)
        # 根据参数采样一百个步骤的x，每隔十步画出来，迭代了4000个周期，逐渐更接近于原始
        x_seq = p_sample_loop(model, dataset.shape, num_steps, betas, one_minus_alphas_bar_sqrt)

        fig, axs = plt.subplots(1, 10, figsize=(28, 3))
        for i in range(1, 11):
            cur_x = x_seq[i * 10].detach()
            axs[i - 1].scatter(cur_x[:, 0], cur_x[:, 1], color='red', edgecolor='white')
            axs[i - 1].set_axis_off()  # 关闭轴线
            axs[i - 1].set_title('$q(\mathbf{x}_{' + str(i * 10) + '})$')
plt.show()

# 前向后向过程gif
imgs = []
for i in range(100):
    plt.clf()
    q_i = forawrd_diffusion(dataset, torch.tensor([i]))  # 获取正向扩散第i个步长时的采样值
    plt.scatter(q_i[:, 0], q_i[:, 1], color='red', edgecolor='white', s=5)
    plt.axis('off')

    # 将采样图像保存PNG再读取,转换为字节数组
    img_buf = io.BytesIO()
    plt.savefig(img_buf, format='png')
    img = Image.open(img_buf)
    imgs.append(img)
plt.show()
reverse = []
for i in range(100):
    plt.clf()
    cur_x = x_seq[i].detach()  # 获取逆向阶段第i个采样值
    plt.scatter(cur_x[:, 0], cur_x[:, 1], color='red', edgecolor='white', s=5)
    plt.axis('off')

    img_buf = io.BytesIO()
    plt.savefig(img_buf, format='png')
    img = Image.open(img_buf)
    reverse.append(img)
plt.show()
imgs = imgs + reverse
imgs[0].save("diffusion.gif", format='GIF', append_images=imgs, save_all=True, duration=100, loop=0)  # 保存为gif
